import numpy as np
import matplotlib.pyplot as plt

fig,ax = plt.subplots()

data = np.loadtxt('SQETOTAL')

sqe = data[:,2].reshape((4096,20),order='F')
sqe[sqe<=0] = 1e-7

np.savetxt('E.txt',data[0:4096,0])
np.savetxt('0p2.txt',sqe[:,3])
np.savetxt('0p4.txt',sqe[:,7])
np.savetxt('0p8.txt',sqe[:,15])

vmin=-1.5
vmax=0.5
fig = ax.imshow(np.log10(sqe),origin='lower',vmin=vmin,vmax=vmax,aspect='auto',cmap='viridis',interpolation='none')


xpos = np.linspace(-1,19,2)
xticks = np.linspace(0,1,2,dtype='int')
ax.set_xticks(xpos)
ax.set_xticklabels(xticks)
ax.set_xlabel('0 0 L (r.l.u)',fontsize=20)

ypos = np.linspace(0,398,6)  # 398 20meV
yticks = np.linspace(0,20,6,dtype='int')
ax.set_yticks(ypos)
ax.set_yticklabels(yticks)
ax.set_ylim(0,398)
ax.set_ylabel('Energy (meV)',fontsize=20)

ax.tick_params(which='both',labelsize=20)

cbar = plt.colorbar(fig)
cbar.ax.tick_params(labelsize=20)

plt.savefig('sqe.png',bbox_inches='tight')
plt.show()
